package com.yjz.middleware.RocketMq;

import lombok.extern.slf4j.Slf4j;
import org.apache.rocketmq.client.exception.MQClientException;
import org.apache.rocketmq.client.producer.LocalTransactionState;
import org.apache.rocketmq.client.producer.SendResult;
import org.apache.rocketmq.client.producer.TransactionListener;
import org.apache.rocketmq.client.producer.TransactionMQProducer;
import org.apache.rocketmq.common.message.Message;
import org.apache.rocketmq.common.message.MessageExt;
import org.apache.rocketmq.remoting.common.RemotingHelper;
import org.junit.Test;

import java.io.UnsupportedEncodingException;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @author :yuanJinZhong
 * @description:事务消息生产者
 * @date :2018/8/7-16:36
 */
@Slf4j
public class TransactionProducerUnitTest {


  @Test
  public void transactionTest() throws MQClientException, UnsupportedEncodingException, InterruptedException {
    TransactionMQProducer tsProducer = new TransactionMQProducer("my_transaction_producer");

    ExecutorService executorService = new ThreadPoolExecutor(2, 5, 100, TimeUnit.SECONDS, new ArrayBlockingQueue<Runnable>(2000), new ThreadFactory() {
      @Override
      public Thread newThread(Runnable r) {
        Thread thread = new Thread(r);
        thread.setName("client-transaction-msg-check-thread");
        return thread;
      }
    });
    tsProducer.setNamesrvAddr("47.98.237.214:9876");
    tsProducer.setExecutorService(executorService);

    tsProducer.setTransactionListener(new TransactionListener() {
      private AtomicInteger transactionIndex = new AtomicInteger(0);
      private ConcurrentHashMap<String, Integer> localTrans = new ConcurrentHashMap<>();

      @Override
      /**
       * When send transactional prepare(half) message succeed, this method will be invoked to execute local transaction.
       */
      public LocalTransactionState executeLocalTransaction(Message msg, Object arg) {
        int value = transactionIndex.getAndIncrement();
        int status = value % 3;
        localTrans.put(msg.getTransactionId(), status);
        return LocalTransactionState.UNKNOW;

      }

      @Override
      public LocalTransactionState checkLocalTransaction(MessageExt msg) {
        Integer status = localTrans.get(msg.getTransactionId());
        if (null != status) {
          switch (status) {
            case 0:
              return LocalTransactionState.UNKNOW;
            case 1:
              return LocalTransactionState.COMMIT_MESSAGE;
            case 2:
              return LocalTransactionState.ROLLBACK_MESSAGE;
          }
        }
        return LocalTransactionState.COMMIT_MESSAGE;
      }
    });
    tsProducer.start();
    String[] tags = new String[]{"TagA", "TagB", "TagC", "TagD", "TagE"};
    for (int i = 0; i < 10; i++) {
      Message msg = new Message("TopicTest1234", tags[i % tags.length], "KEY" + i,
        ("Hello RocketMQ " + i).getBytes(RemotingHelper.DEFAULT_CHARSET));
      SendResult result = tsProducer.sendMessageInTransaction(msg, null);
      log.info("{}", result);
      Thread.sleep(10);
    }
    tsProducer.shutdown();
  }

}
